#!/usr/bin/env python
'''
This function extracts the sequences in the given junction regions.
The 1st arg is the sequence file *.fa
The 2nd arg is the junctions file with format
The 3rd arg is the output file, which is the sequences. 
chrName, start, end, geneName, support, Strand
'''
__author__ = "Yangho Chen"
__email__ = "yanghoch@usc.edu"
__modification__ = "Created on Aug 16, 2009"

import string
import sys
import os
import getopt
from TophatOutFormat import JunctionLineInfo
from parse_fasta_file import get_tag_seq_from_fasta_w_substr
from parse_fasta_file import reverseCompliment

def get_juncs_seq_from_fa( seqFileN, juncFileN, oFileN ):
    def assertJunction( leftS, rightS, halfJuncL ):
        ''' Print the pattern of the splicing junction region.
            leftS is the left starting position of junction region.
            rightS is the right starting position of the junction region.
            halfJuncL is the length of junction region.   
        '''
        juncL = 2 * halfJuncL
        leftE = seq[ leftS : leftS + halfJuncL].upper()
        leftI = seq[ leftS + halfJuncL : leftS + juncL].lower()
        rightI = seq[ rightS: rightS + halfJuncL].lower()
        rightE = seq[ rightS + halfJuncL : rightS + juncL].upper()
        if junc.strand == '+':
            exon2intron = leftE + leftI
            intron2exon = rightI + rightE
        else:
            exon2intron = string.join( reverseCompliment( rightI + rightE ), '' )
            intron2exon = string.join( reverseCompliment( leftE + leftI ), '' )
        if len( exon2intron ) < juncL or len( intron2exon ) < juncL:
            print seqFileN, "has length", len( seq ), '<', rightSeqStart
            return False
        elif( exon2intron[halfJuncL] != 'g' or \
              exon2intron[halfJuncL + 1] != 't' or \
              intron2exon[halfJuncL - 2] != 'a' or \
              intron2exon[halfJuncL - 1] != 'g' ):
                if junc.strand == '+':
                    print "Special Junction seq", junc.strand, junc.noMaps, \
                            leftS, exon2intron, rightS, intron2exon
                else:
                    print "Special Junction seq", junc.strand, junc.noMaps, \
                            rightS, exon2intron, leftS, intron2exon
                return False
        return True

    faFile = open( seqFileN, "r" )
    iFile = open( juncFileN, "r" )
    oFile = open( oFileN, "w" )
    FLANKING_REGION_LENGTH = 5 # The length of extracted seq in each side
    [tag, seq] = get_tag_seq_from_fasta_w_substr( faFile, ">" )
    correctJucNo = 0
    for line in iFile:
        junc = JunctionLineInfo( line )
        # Assume the given junction are 1-coordinate, translate to 0-coordinate
        leftSeqStart = junc.start - FLANKING_REGION_LENGTH + 1 - 1
        rightSeqStart = junc.end + 1 - 1 # Junction end is the end of intron
        juncSeq = seq[ leftSeqStart : leftSeqStart + FLANKING_REGION_LENGTH ]\
                + seq[ rightSeqStart : rightSeqStart + FLANKING_REGION_LENGTH ]
        if junc.strand != '+':
            juncSeq = string.join( reverseCompliment( juncSeq ), '' )
        if assertJunction( leftSeqStart, rightSeqStart - FLANKING_REGION_LENGTH\
                          , FLANKING_REGION_LENGTH ):
            correctJucNo += 1
            #print juncSeq
        oFile.write( juncSeq + '\n' ) # record both normal and special junctions
    faFile.close();
    iFile.close()
    oFile.close()
    print "There are ", correctJucNo, "normal junctions"

def usage():
    sys.stderr.write( "Working dir is " + os.getcwd() + '\n' )
    print __doc__

def main():
    try:
        opts, args = getopt.getopt( sys.argv[1:], 'h', ["help"] )
    except getopt.error, msg:
        print msg
        sys.exit( 2 );
    #process options and argument 
    for o in opts:
        if o in ( "-h", "--help" ):
            print __doc__
            sys.exit();
    if ( len( args ) == 3 ):
        seqFileN = args[0]
        juncFileN = args[1]
        oFileN = args[2]
        if( not os.path.isfile( seqFileN ) ):
            print "File ", seqFileN, " are not found."
        elif ( not os.path.isfile( juncFileN ) ):
            print "File ", juncFileN, " are not found."
        else:
            get_juncs_seq_from_fa( seqFileN, juncFileN, oFileN )
    else:
        print __doc__

if __name__ == '__main__':
    main()
